#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2; 

//bf16
template<typename typeT> class KernelLinSpaceBF16 {
public:
    __aicore__ inline KernelLinSpaceBF16() {}
    __aicore__ inline void Init(GM_ADDR start, GM_ADDR stop, GM_ADDR num_axes, GM_ADDR output,
                                uint32_t CoreDataNum,
                                uint32_t finalTileNum,
                                uint32_t tileDataNum,
                                uint32_t TailDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->CoreDataNum = CoreDataNum;
        this->finalTileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->TailDataNum = TailDataNum;

        startGm.SetGlobalBuffer((__gm__ typeT*)start, 1);
        stopGm.SetGlobalBuffer((__gm__ typeT*)stop, 1);
        num_axesGm.SetGlobalBuffer((__gm__ int32_t*)num_axes, 1);

        this->num_axes = num_axesGm.GetValue(0);
        outputGm.SetGlobalBuffer((__gm__ typeT*)output, this->num_axes);
        pipe.InitBuffer(tmp1, 16 * sizeof(float));
        pipe.InitBuffer(tmp2, 16 * sizeof(typeT));
        // pipe alloc memory to queue, the unit is Bytes
        // pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileDataNum * sizeof(typeT));
        // pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileDataNum * sizeof(typeT));
        // // pipe.InitBuffer(tmp1, this->tileDataNum * sizeof(float));
        // pipe.InitBuffer(tmp2, this->tileDataNum * sizeof(int64_t));
        // pipe.InitBuffer(tmp3, this->tilelength * sizeof(int8_t) / 8);
    }
    __aicore__ inline void Process() {
        LocalTensor<float> tmp1Local = tmp1.Get<float>();
        LocalTensor<typeT> tmp2Local = tmp2.Get<typeT>();
        if constexpr (std::is_same_v<typeT, __bf16>) {
            typeT start = startGm.GetValue(0);
            typeT stop = stopGm.GetValue(0);
            outputGm.SetValue(0, start);
            outputGm.SetValue(this->num_axes - 1, stop);
            tmp2Local.SetValue(0, start);
            tmp2Local.SetValue(1, stop);
            Cast(tmp1Local, tmp2Local, RoundMode::CAST_NONE, 16);
            float startfp = tmp1Local.GetValue(0);
            float stopfp = tmp1Local.GetValue(1);
            float interval = (stopfp - startfp) / (this->num_axes - 1);
            for(int32_t k = 1; k < this->num_axes - 1; k++) {
                float tmpvalue = startfp + (k * interval);
                tmp1Local.SetValue(0, tmpvalue);
                Cast(tmp2Local, tmp1Local, RoundMode::CAST_RINT, 16);
                typeT tmpbf16 = tmp2Local.GetValue(0);
                outputGm.SetValue(k, tmpbf16);
            }
        }
      
    }

private:
    TPipe pipe;
    TBuf<QuePosition::VECCALC> tmp1, tmp2;
    GlobalTensor<typeT> startGm;
    GlobalTensor<typeT> stopGm;
    GlobalTensor<int32_t> num_axesGm;
    GlobalTensor<typeT> outputGm;
    uint32_t CoreDataNum;
    uint32_t finalTileNum;
    uint32_t tileDataNum;
    uint32_t TailDataNum;
    int32_t num_axes;
};
//INT8
template<typename typeT> class KernelLinSpaceINT8 {
public:
    __aicore__ inline KernelLinSpaceINT8() {}
    __aicore__ inline void Init(GM_ADDR start, GM_ADDR stop, GM_ADDR num_axes, GM_ADDR output,
                                uint32_t CoreDataNum,
                                uint32_t finalTileNum,
                                uint32_t tileDataNum,
                                uint32_t TailDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->CoreDataNum = CoreDataNum;
        this->finalTileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->TailDataNum = TailDataNum;

        startGm.SetGlobalBuffer((__gm__ typeT*)start, 1);
        stopGm.SetGlobalBuffer((__gm__ typeT*)stop, 1);
        num_axesGm.SetGlobalBuffer((__gm__ int32_t*)num_axes, 1);

        this->num_axes = num_axesGm.GetValue(0);
        outputGm.SetGlobalBuffer((__gm__ typeT*)output, this->num_axes);
        pipe.InitBuffer(tmp1, 16 * sizeof(float));
        pipe.InitBuffer(tmp2, 16 * sizeof(typeT));
        pipe.InitBuffer(tmp3, 16 * sizeof(half));

    }
    __aicore__ inline void Process() {
        LocalTensor<float> tmp1Local = tmp1.Get<float>();
        LocalTensor<typeT> tmp2Local = tmp2.Get<typeT>();
        LocalTensor<half> tmp3Local = tmp3.Get<half>();
        if constexpr (std::is_same_v<typeT, int8_t>) {
            float start = startGm.GetValue(0);
            float stop = stopGm.GetValue(0);
            outputGm.SetValue(0, start);
            outputGm.SetValue(this->num_axes - 1, stop);
            float interval = (stop - start) / (this->num_axes - 1);
            half tmpfp16 = interval;
            interval = tmpfp16;
            for(int32_t k = 1; k < this->num_axes - 1; k++) {
                tmp3Local.SetValue(0, interval);
                Muls(tmp3Local, tmp3Local, half(k), 1);
                Adds(tmp3Local, tmp3Local, half(start), 1);
                //float tmpvalue = start + (k * interval);
                //half tmpfp16 = tmpvalue;
                //tmp3Local.SetValue(0, tmpfp16);
                Cast(tmp2Local, tmp3Local, RoundMode::CAST_TRUNC, 16);
                typeT tmpint = tmp2Local.GetValue(0);
                outputGm.SetValue(k, tmpint);
            }
            
        }
      
    }

private:
    TPipe pipe;
    TBuf<QuePosition::VECCALC> tmp1, tmp2, tmp3;
    GlobalTensor<typeT> startGm;
    GlobalTensor<typeT> stopGm;
    GlobalTensor<int32_t> num_axesGm;
    GlobalTensor<typeT> outputGm;
    uint32_t CoreDataNum;
    uint32_t finalTileNum;
    uint32_t tileDataNum;
    uint32_t TailDataNum;
    int32_t num_axes;
};
//UINT8
template<typename typeT> class KernelLinSpaceUINT8 {
public:
    __aicore__ inline KernelLinSpaceUINT8() {}
    __aicore__ inline void Init(GM_ADDR start, GM_ADDR stop, GM_ADDR num_axes, GM_ADDR output,
                                uint32_t CoreDataNum,
                                uint32_t finalTileNum,
                                uint32_t tileDataNum,
                                uint32_t TailDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->CoreDataNum = CoreDataNum;
        this->finalTileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->TailDataNum = TailDataNum;

        startGm.SetGlobalBuffer((__gm__ typeT*)start, 1);
        stopGm.SetGlobalBuffer((__gm__ typeT*)stop, 1);
        num_axesGm.SetGlobalBuffer((__gm__ int32_t*)num_axes, 1);

        this->num_axes = num_axesGm.GetValue(0);
        outputGm.SetGlobalBuffer((__gm__ typeT*)output, this->num_axes);
        pipe.InitBuffer(tmp1, 16 * sizeof(float));
        pipe.InitBuffer(tmp2, 16 * sizeof(typeT));
        pipe.InitBuffer(tmp3, 16 * sizeof(half));

    }
    __aicore__ inline void Process() {
        LocalTensor<float> tmp1Local = tmp1.Get<float>();
        LocalTensor<typeT> tmp2Local = tmp2.Get<typeT>();
        LocalTensor<half> tmp3Local = tmp3.Get<half>();
        if constexpr (std::is_same_v<typeT, uint8_t>) {
            typeT startori = startGm.GetValue(0);
            typeT stopori = stopGm.GetValue(0);
            tmp2Local.SetValue(0, startori);
            tmp2Local.SetValue(1, stopori);
            Cast(tmp3Local, tmp2Local, RoundMode::CAST_NONE, 16);
            float start = tmp3Local.GetValue(0);
            float stop = tmp3Local.GetValue(1);
            outputGm.SetValue(0, startori);
            outputGm.SetValue(this->num_axes - 1, stopori);
            float interval = (stop - start) / (this->num_axes - 1);
            for(int32_t k = 1; k < this->num_axes - 1; k++) {
                float tmpvalue = start + (k * interval);
                tmp1Local.SetValue(0, tmpvalue);
                Cast(tmp3Local, tmp1Local, RoundMode::CAST_RINT, 16);
                Cast(tmp2Local, tmp3Local, RoundMode::CAST_TRUNC, 16);
                typeT tmpint = tmp2Local.GetValue(0);
                outputGm.SetValue(k, tmpint);
            }
        }
      
    }

private:
    TPipe pipe;
    TBuf<QuePosition::VECCALC> tmp1, tmp2, tmp3;
    GlobalTensor<typeT> startGm;
    GlobalTensor<typeT> stopGm;
    GlobalTensor<int32_t> num_axesGm;
    GlobalTensor<typeT> outputGm;
    uint32_t CoreDataNum;
    uint32_t finalTileNum;
    uint32_t tileDataNum;
    uint32_t TailDataNum;
    int32_t num_axes;
};
//INT16
template<typename typeT> class KernelLinSpaceINT16 {
public:
    __aicore__ inline KernelLinSpaceINT16() {}
    __aicore__ inline void Init(GM_ADDR start, GM_ADDR stop, GM_ADDR num_axes, GM_ADDR output,
                                uint32_t CoreDataNum,
                                uint32_t finalTileNum,
                                uint32_t tileDataNum,
                                uint32_t TailDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->CoreDataNum = CoreDataNum;
        this->finalTileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->TailDataNum = TailDataNum;

        startGm.SetGlobalBuffer((__gm__ typeT*)start, 1);
        stopGm.SetGlobalBuffer((__gm__ typeT*)stop, 1);
        num_axesGm.SetGlobalBuffer((__gm__ int32_t*)num_axes, 1);

        this->num_axes = num_axesGm.GetValue(0);
        outputGm.SetGlobalBuffer((__gm__ typeT*)output, this->num_axes);
        pipe.InitBuffer(tmp1, 16 * sizeof(float));
        pipe.InitBuffer(tmp2, 16 * sizeof(typeT));

    }
    __aicore__ inline void Process() {
        LocalTensor<float> tmp1Local = tmp1.Get<float>();
        LocalTensor<typeT> tmp2Local = tmp2.Get<typeT>();
        if constexpr (std::is_same_v<typeT, int16_t>) {
            float start = startGm.GetValue(0);
            float stop = stopGm.GetValue(0);
            outputGm.SetValue(0, start);
            outputGm.SetValue(this->num_axes - 1, stop);
            float interval = (stop - start) / (this->num_axes - 1);
            for(int32_t k = 1; k < this->num_axes - 1; k++) {
                float tmpvalue = start + (k * interval);
                tmp1Local.SetValue(0, tmpvalue);
                Cast(tmp2Local, tmp1Local, RoundMode::CAST_TRUNC, 16);
                typeT tmpint = tmp2Local.GetValue(0);
                outputGm.SetValue(k, tmpint);
            }
        }
      
    }

private:
    TPipe pipe;
    TBuf<QuePosition::VECCALC> tmp1, tmp2;
    GlobalTensor<typeT> startGm;
    GlobalTensor<typeT> stopGm;
    GlobalTensor<int32_t> num_axesGm;
    GlobalTensor<typeT> outputGm;
    uint32_t CoreDataNum;
    uint32_t finalTileNum;
    uint32_t tileDataNum;
    uint32_t TailDataNum;
    int32_t num_axes;
};

//INT32
template<typename typeT> class KernelLinSpaceINT32 {
public:
    __aicore__ inline KernelLinSpaceINT32() {}
    __aicore__ inline void Init(GM_ADDR start, GM_ADDR stop, GM_ADDR num_axes, GM_ADDR output,
                                uint32_t CoreDataNum,
                                uint32_t finalTileNum,
                                uint32_t tileDataNum,
                                uint32_t TailDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->CoreDataNum = CoreDataNum;
        this->finalTileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->TailDataNum = TailDataNum;

        startGm.SetGlobalBuffer((__gm__ typeT*)start, 1);
        stopGm.SetGlobalBuffer((__gm__ typeT*)stop, 1);
        num_axesGm.SetGlobalBuffer((__gm__ int32_t*)num_axes, 1);

        this->num_axes = num_axesGm.GetValue(0);
        outputGm.SetGlobalBuffer((__gm__ typeT*)output, this->num_axes);
        pipe.InitBuffer(tmp1, 16 * sizeof(float));
        pipe.InitBuffer(tmp2, 16 * sizeof(typeT));

    }
    __aicore__ inline void Process() {
        LocalTensor<float> tmp1Local = tmp1.Get<float>();
        LocalTensor<typeT> tmp2Local = tmp2.Get<typeT>();
        if constexpr (std::is_same_v<typeT, int32_t>) {
            float start = startGm.GetValue(0);
            float stop = stopGm.GetValue(0);
            outputGm.SetValue(0, start);
            outputGm.SetValue(this->num_axes - 1, stop);
            float interval = (stop - start) / (this->num_axes - 1);
            for(int32_t k = 1; k < this->num_axes - 1; k++) {
                float tmpvalue = start + (k * interval);
                tmp1Local.SetValue(0, tmpvalue);
                Cast(tmp2Local, tmp1Local, RoundMode::CAST_TRUNC, 16);
                typeT tmpint = tmp2Local.GetValue(0);
                outputGm.SetValue(k, tmpint);
            }
        }
      
    }

private:
    TPipe pipe;
    TBuf<QuePosition::VECCALC> tmp1, tmp2;
    GlobalTensor<typeT> startGm;
    GlobalTensor<typeT> stopGm;
    GlobalTensor<int32_t> num_axesGm;
    GlobalTensor<typeT> outputGm;
    uint32_t CoreDataNum;
    uint32_t finalTileNum;
    uint32_t tileDataNum;
    uint32_t TailDataNum;
    int32_t num_axes;
};


//HALF
template<typename typeT> class KernelLinSpaceHALF {
public:
    __aicore__ inline KernelLinSpaceHALF() {}
    __aicore__ inline void Init(GM_ADDR start, GM_ADDR stop, GM_ADDR num_axes, GM_ADDR output,
                                uint32_t CoreDataNum,
                                uint32_t finalTileNum,
                                uint32_t tileDataNum,
                                uint32_t TailDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->CoreDataNum = CoreDataNum;
        this->finalTileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->TailDataNum = TailDataNum;

        startGm.SetGlobalBuffer((__gm__ typeT*)start, 1);
        stopGm.SetGlobalBuffer((__gm__ typeT*)stop, 1);
        num_axesGm.SetGlobalBuffer((__gm__ int32_t*)num_axes, 1);

        this->num_axes = num_axesGm.GetValue(0);
        outputGm.SetGlobalBuffer((__gm__ typeT*)output, this->num_axes);

        // pipe alloc memory to queue, the unit is Bytes
        // pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileDataNum * sizeof(typeT));
        // pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileDataNum * sizeof(typeT));
        // // pipe.InitBuffer(tmp1, this->tileDataNum * sizeof(float));
        // pipe.InitBuffer(tmp2, this->tileDataNum * sizeof(int64_t));
        // pipe.InitBuffer(tmp3, this->tilelength * sizeof(int8_t) / 8);
    }
    __aicore__ inline void Process() {
        if constexpr (std::is_same_v<typeT, half>) {
            float start = startGm.GetValue(0);
            float stop = stopGm.GetValue(0);
            outputGm.SetValue(0, start);
            outputGm.SetValue(this->num_axes - 1, stop);
            float interval = (stop - start) / (this->num_axes - 1);
            for(int32_t k = 1; k < this->num_axes - 1; k++) {
                float tmpvalue = start + (k * interval);
                outputGm.SetValue(k, tmpvalue);
            }
        }
      
    }

private:
    TPipe pipe;

    GlobalTensor<typeT> startGm;
    GlobalTensor<typeT> stopGm;
    GlobalTensor<int32_t> num_axesGm;
    GlobalTensor<typeT> outputGm;
    uint32_t CoreDataNum;
    uint32_t finalTileNum;
    uint32_t tileDataNum;
    uint32_t TailDataNum;
    int32_t num_axes;
};


//FLOAT32
template<typename typeT> class KernelLinSpaceFLOAT {
public:
    __aicore__ inline KernelLinSpaceFLOAT() {}
    __aicore__ inline void Init(GM_ADDR start, GM_ADDR stop, GM_ADDR num_axes, GM_ADDR output,
                                uint32_t CoreDataNum,
                                uint32_t finalTileNum,
                                uint32_t tileDataNum,
                                uint32_t TailDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->CoreDataNum = CoreDataNum;
        this->finalTileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->TailDataNum = TailDataNum;

        startGm.SetGlobalBuffer((__gm__ typeT*)start, 1);
        stopGm.SetGlobalBuffer((__gm__ typeT*)stop, 1);
        num_axesGm.SetGlobalBuffer((__gm__ int32_t*)num_axes, 1);

        this->num_axes = num_axesGm.GetValue(0);
        outputGm.SetGlobalBuffer((__gm__ typeT*)output, this->num_axes);

        // pipe alloc memory to queue, the unit is Bytes
        // pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileDataNum * sizeof(typeT));
        // pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileDataNum * sizeof(typeT));
        // // pipe.InitBuffer(tmp1, this->tileDataNum * sizeof(float));
        // pipe.InitBuffer(tmp2, this->tileDataNum * sizeof(int64_t));
        // pipe.InitBuffer(tmp3, this->tilelength * sizeof(int8_t) / 8);
    }
    __aicore__ inline void Process() {
        if constexpr (std::is_same_v<typeT, float>) {
            float start = startGm.GetValue(0);
            float stop = stopGm.GetValue(0);
            outputGm.SetValue(0, start);
            outputGm.SetValue(this->num_axes - 1, stop);
            float interval = (stop - start) / (this->num_axes - 1);
            for(int32_t k = 1; k < this->num_axes - 1; k++) {
                float tmpvalue = start + (k * interval);
                outputGm.SetValue(k, tmpvalue);
            }
        }
      
    }

private:
    TPipe pipe;

    GlobalTensor<typeT> startGm;
    GlobalTensor<typeT> stopGm;
    GlobalTensor<int32_t> num_axesGm;
    GlobalTensor<typeT> outputGm;
    uint32_t CoreDataNum;
    uint32_t finalTileNum;
    uint32_t tileDataNum;
    uint32_t TailDataNum;
    int32_t num_axes;
};

//Special case1
template<typename typeT> class KernelLinSpaceSP1 {
public:
    __aicore__ inline KernelLinSpaceSP1() {}
    __aicore__ inline void Init(GM_ADDR start, GM_ADDR stop, GM_ADDR num_axes, GM_ADDR output) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        startGm.SetGlobalBuffer((__gm__ typeT*)start, 1);
        stopGm.SetGlobalBuffer((__gm__ typeT*)stop, 1);
        num_axesGm.SetGlobalBuffer((__gm__ int32_t*)num_axes, 1);

        this->num_axes = num_axesGm.GetValue(0);
        outputGm.SetGlobalBuffer((__gm__ typeT*)output, this->num_axes);
        typeT tmp = startGm.GetValue(0);
        outputGm.SetValue(0, tmp);
    }

private:
    TPipe pipe;

    GlobalTensor<typeT> startGm;
    GlobalTensor<typeT> stopGm;
    GlobalTensor<int32_t> num_axesGm;
    GlobalTensor<typeT> outputGm;
    int32_t num_axes;
};

extern "C" __global__ __aicore__ void lin_space(GM_ADDR start, GM_ADDR stop, GM_ADDR num_axes, GM_ADDR output, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    if (TILING_KEY_IS(100)) {
        KernelLinSpaceBF16<DTYPE_START> op;
        op.Init(start, stop, num_axes, output,
            tiling_data.CoreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum);  
        op.Process();   
    } else if (TILING_KEY_IS(200)) {
        //int8
        KernelLinSpaceINT8<DTYPE_START> op;
        op.Init(start, stop, num_axes, output,
            tiling_data.CoreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum);  
        op.Process(); 
    } else if (TILING_KEY_IS(201)) {
        //int32
        KernelLinSpaceINT32<DTYPE_START> op;
        op.Init(start, stop, num_axes, output,
            tiling_data.CoreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum);  
        op.Process();   
    } else if (TILING_KEY_IS(202)) {
        //uint8
        KernelLinSpaceUINT8<DTYPE_START> op;
        op.Init(start, stop, num_axes, output,
            tiling_data.CoreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum);  
        op.Process();            
    } else if (TILING_KEY_IS(203)) {
        //int16
        KernelLinSpaceINT16<DTYPE_START> op;
        op.Init(start, stop, num_axes, output,
            tiling_data.CoreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum);  
        op.Process();            
    } else if (TILING_KEY_IS(300)) {
        KernelLinSpaceHALF<DTYPE_START> op;
        op.Init(start, stop, num_axes, output,
            tiling_data.CoreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum);  
        op.Process();      
    } else if (TILING_KEY_IS(400)) {
        KernelLinSpaceFLOAT<DTYPE_START> op;
        op.Init(start, stop, num_axes, output,
            tiling_data.CoreDataNum, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum);  
        op.Process();
    } else if (TILING_KEY_IS(900)) {
        KernelLinSpaceSP1<DTYPE_START> op;
        op.Init(start, stop, num_axes, output);  
    }    
}